import matplotlib
matplotlib.use('Agg')  # 设置后端为Agg，避免多线程问题

import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
import argparse
import logging
from datetime import datetime, timedelta
from factor_validation import FactorValidator
# 注意：DataCacheManager会在main函数中按需导入

def get_future_returns(self, trade_date, n_days=5):
        """
        计算从trade_date开始未来n_days的行业收益率
        """
def setup_logging():
    """
    设置日志记录配置
    """
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'validation.log')),
            logging.StreamHandler()
        ]
    )

def validate_date(date_str):
    """
    验证日期格式是否正确
    """
    try:
        datetime.strptime(date_str, '%Y%m%d')
        return True
    except ValueError:
        return False

def preload_data(validator, start_date, end_date, n_days):
    """
    预加载验证过程中需要的所有数据
    优化版本：使用并行加载和批处理
    """
    logger = logging.getLogger(__name__)
    logger.info("开始预加载数据...")
    
    # 获取交易日列表
    trade_dates = validator.analyzer.get_trade_dates(start_date, end_date)
    if not trade_dates:
        logger.error(f"无法获取交易日列表: {start_date}-{end_date}")
        return False
    
    logger.info(f"预加载期间共有{len(trade_dates)}个交易日")
    
    # 使用并行方式预加载所有数据
    import concurrent.futures
    import threading
    
    # 创建线程锁
    lock = threading.Lock()
    
    # 定义预加载函数
    def preload_all_data():
        # 预加载股票数据
        validator.cache_manager.preload_stock_data(trade_dates)
        return "股票数据"
        
    def preload_future_returns():
        # 预加载未来收益率数据
        validator.cache_manager.preload_future_returns(trade_dates, n_days)
        return "未来收益率数据"
        
    def preload_index_data():
        # 预加载指数数据
        validator.cache_manager.preload_index_data(trade_dates)
        return "指数数据"
        
    def preload_limit_data():
        # 预加载涨停数据
        validator.cache_manager.preload_limit_data(trade_dates)
        return "涨停数据"
    
    # 并行执行所有预加载任务
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
        futures = [
            executor.submit(preload_all_data),
            executor.submit(preload_future_returns),
            executor.submit(preload_index_data),
            executor.submit(preload_limit_data)
        ]
        
        for future in concurrent.futures.as_completed(futures):
            try:
                data_type = future.result()
                logger.info(f"完成{data_type}预加载")
            except Exception as e:
                logger.error(f"预加载数据失败: {str(e)}")
    
    logger.info("所有数据预加载完成")
    return True

def main():
    # 设置日志
    setup_logging()
    logger = logging.getLogger(__name__)
    
    try:
        # 创建命令行参数解析器
        parser = argparse.ArgumentParser(description='因子有效性验证框架')
        parser.add_argument('--start_date', type=str, default='20230213', help='回测起始日期，格式：YYYYMMDD')
        parser.add_argument('--end_date', type=str, default='20230317', help='回测结束日期，格式：YYYYMMDD')
        parser.add_argument('--n_days', type=int, default=5, help='未来收益计算天数')
        parser.add_argument('--factors', type=str, nargs='+', 
                          default=['density', 'consecutive', 'historical', 'strength'],
                          help='需要验证的因子列表')
        parser.add_argument('--clear_cache', action='store_true', help='清除数据缓存')
        parser.add_argument('--cache_days', type=int, default=None, help='清除指定天数之前的缓存数据')
        parser.add_argument('--no_preload', action='store_true', help='禁用数据预加载')
        # 添加并行处理参数
        parser.add_argument('--parallel', action='store_true', help='启用并行处理多个因子')
        parser.add_argument('--max_workers', type=int, default=None, help='最大并行工作线程数，默认为CPU核心数')
        args = parser.parse_args()
        
        # 验证日期格式
        if not validate_date(args.start_date) or not validate_date(args.end_date):
            logger.error("日期格式错误，请使用YYYYMMDD格式")
            return
            
        if args.start_date > args.end_date:
            logger.error("起始日期不能晚于结束日期")
            return
    
    except Exception as e:
        logger.error(f"参数解析失败: {str(e)}")
        return
    
    # 导入数据缓存管理器
    from data_cache_manager import DataCacheManager
    
    # 创建数据缓存目录
    cache_manager = DataCacheManager()
    logger.info(f"数据缓存目录: {cache_manager.cache_dir}")
    
    # 处理缓存清理请求
    if args.clear_cache:
        count = cache_manager.clear_cache(args.cache_days)
        logger.info(f"已清除{count}个缓存文件")
        if not args.factors:  # 如果只是清除缓存而不运行验证
            return
            
    # 初始化因子验证器
    validator = FactorValidator()
    
    # 创建输出目录
    output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'output')
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        
    # 预加载数据
    if not args.no_preload:
        preload_success = preload_data(validator, args.start_date, args.end_date, args.n_days)
        if not preload_success:
            logger.warning("数据预加载失败，将在验证过程中按需加载数据")
    else:
        logger.info("已禁用数据预加载，将在验证过程中按需加载数据")
    
    # 导入多线程模块
    import concurrent.futures
    import threading
    
    # 创建线程锁，用于保护matplotlib绘图
    plot_lock = threading.Lock()
    
    # 定义单个因子验证的函数
    def validate_single_factor(factor_name, start_date, end_date, n_days, output_dir, shared_data=None):
        thread_logger = logging.getLogger(f"{__name__}.{factor_name}")
        thread_logger.info(f"\n{'='*50}")
        thread_logger.info(f"开始验证因子: {factor_name}")
        thread_logger.info(f"{'='*50}")
        
        try:
            # 运行因子验证，并获取加载的数据以便后续步骤重用
            factor_data = validator.validate_factor(factor_name, start_date, end_date, n_days, shared_data=shared_data)
            
            # 运行市场状态测试
            print(f"\n4. {factor_name} 市场状态测试")
            try:
                # 使用已加载的数据进行市场状态测试
                market_state_ic = validator.run_market_state_test(factor_name, start_date, end_date, n_days, all_data=factor_data)
                if market_state_ic and not all(np.isnan(v) for v in market_state_ic.values()):
                    print(f"{factor_name} 不同市场状态下的平均IC值:")
                    for state, ic in market_state_ic.items():
                        print(f"{state}市场: {ic:.4f}")
                
                # 使用线程锁保护matplotlib绘图
                with plot_lock:
                    # 绘制市场状态测试图
                    plt.figure(figsize=(8, 5))
                    states = list(market_state_ic.keys())
                    ics = [market_state_ic[s] for s in states]
                    plt.bar(range(len(states)), ics)
                    plt.xticks(range(len(states)), states)
                    plt.title(f'{factor_name} Performance in Different Market States')
                    plt.xlabel('Market State')
                    plt.ylabel('Average IC')
                    plt.grid(True)
                    plt.tight_layout()
                    
                    # 保存图表
                    plt.savefig(os.path.join(output_dir, f'{factor_name}_market_state.png'))
                    plt.close()
            except Exception as e:
                print(f"{factor_name} 无法进行有效的市场状态测试: {str(e)}")
                
            # 运行风险调整收益分析
            print(f"\n5. {factor_name} 风险调整收益分析")
            try:
                # 使用已加载的数据进行风险调整收益分析
                risk_results = validator.run_risk_adjusted_returns(factor_name, start_date, end_date, n_days, all_data=factor_data)
                if risk_results and 'high_cumulative' in risk_results and len(risk_results['high_cumulative']) > 0:
                    print(f"{factor_name} 最大回撤: {risk_results['max_drawdown']:.4f}")
                    print(f"{factor_name} 夏普比率: {risk_results['sharpe_ratio']:.4f}")
                    print(f"{factor_name} Alpha: {risk_results['alpha']:.4f}")
                    print(f"{factor_name} Beta: {risk_results['beta']:.4f}")
                
                    # 使用线程锁保护matplotlib绘图
                    with plot_lock:
                        # 绘制高分组、低分组和市场的累积收益率对比图
                        plt.figure(figsize=(10, 6))
                        plt.plot(risk_results['high_cumulative'], label='High Group')
                        plt.plot(risk_results['low_cumulative'], label='Low Group')
                        plt.plot(risk_results['market_cumulative'], label='Market')
                        plt.title(f'{factor_name} Risk-Adjusted Return Analysis')
                        plt.xlabel('Trading Day Index')
                        plt.ylabel('Cumulative Return')
                        plt.legend()
                        plt.grid(True)
                        plt.tight_layout()
                        
                        # 保存图表
                        plt.savefig(os.path.join(output_dir, f'{factor_name}_risk_adjusted.png'))
                        plt.close()
                else:
                    print(f"{factor_name} 无法获取有效的风险调整收益数据")
            except Exception as e:
                thread_logger.error(f"{factor_name} 风险调整收益分析失败: {str(e)}", exc_info=True)
            
            thread_logger.info(f"\n{'='*50}")
            thread_logger.info(f"因子 {factor_name} 验证完成")
            thread_logger.info(f"{'='*50}\n")
            return True, factor_data
        except Exception as e:
            thread_logger.error(f"{factor_name} 验证失败: {str(e)}", exc_info=True)
            return False, None
    
    # 创建共享数据字典，用于存储已加载的数据
    shared_data = {}
    
    # 使用并行处理验证多个因子
    if args.parallel and len(args.factors) > 1:
        logger.info(f"启用并行处理，同时验证{len(args.factors)}个因子")
        max_workers = args.max_workers or None  # None表示使用默认值（CPU核心数）
        
        # 首先验证第一个因子，加载所有数据
        logger.info(f"首先验证第一个因子 {args.factors[0]}，加载所有必要数据")
        success, first_factor_data = validate_single_factor(args.factors[0], args.start_date, args.end_date, args.n_days, output_dir)
        if success and first_factor_data:
            shared_data = first_factor_data
            logger.info(f"成功加载共享数据，包含 {len(shared_data.get('factor_scores', {}))} 个交易日的数据")
        
        # 并行验证剩余因子，共享已加载的数据
        remaining_factors = args.factors[1:]
        if remaining_factors:
            logger.info(f"并行验证剩余的 {len(remaining_factors)} 个因子")
            with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
                # 提交所有因子验证任务，传递共享数据
                future_to_factor = {executor.submit(validate_single_factor, factor, args.start_date, args.end_date, args.n_days, output_dir, shared_data): factor for factor in remaining_factors}
                
                # 获取结果
                for future in concurrent.futures.as_completed(future_to_factor):
                    factor = future_to_factor[future]
                    try:
                        success, _ = future.result()
                        if success:
                            logger.info(f"因子 {factor} 验证任务已完成")
                        else:
                            logger.error(f"因子 {factor} 验证任务失败")
                    except Exception as e:
                        logger.error(f"因子 {factor} 验证任务发生异常: {str(e)}", exc_info=True)
    else:
        # 顺序验证每个因子，共享数据
        for i, factor_name in enumerate(args.factors):
            if i == 0:
                # 第一个因子加载所有数据
                success, factor_data = validate_single_factor(factor_name, args.start_date, args.end_date, args.n_days, output_dir)
                if success and factor_data:
                    shared_data = factor_data
                    logger.info(f"成功加载共享数据，包含 {len(shared_data.get('factor_scores', {}))} 个交易日的数据")
            else:
                # 后续因子使用共享数据
                validate_single_factor(factor_name, args.start_date, args.end_date, args.n_days, output_dir, shared_data)
            
    logger.info(f"\n{'='*50}")
    logger.info(f"所有因子验证完成")
    logger.info(f"{'='*50}\n")
    
    logger.info("\n所有因子验证完成，结果已保存至output目录")

if __name__ == "__main__":
    main()